import networkx as nx
import matplotlib.pyplot as plt


class Node:
    def __init__(self, id):
        self.id = id
        self.neighbors = []
        self.routing_table = {}  # origin->prev_hop for RREQ reverse path
        self.received_rreq = set()  # origins for which RREQ has been processed


# Define network topology
adjacency = {
    1: [2, 3],
    2: [1, 3, 4, 5],
    3: [1, 2, 5],
    4: [2, 6],
    5: [2, 3, 6],
    6: [4, 5]
}

# Initialize nodes
nodes = {nid: Node(nid) for nid in adjacency}
for nid, node in nodes.items():
    node.neighbors = adjacency[nid]

# Get user input with validation
valid_nodes = list(nodes.keys())
print("可用节点:", valid_nodes)

while True:
    try:
        source = int(input("请输入源节点: "))
        if source in valid_nodes:
            break
        print("无效节点!")
    except ValueError:
        print("请输入数字!")

while True:
    try:
        destination = int(input("请输入目的节点: "))
        if destination in valid_nodes and destination != source:
            break
        print("无效节点或与源节点相同!")
    except ValueError:
        print("请输入数字!")

# Build NetworkX graph
G = nx.Graph()
for u in adjacency:
    G.add_node(u)
    for v in adjacency[u]:
        if not G.has_edge(u, v):
            G.add_edge(u, v)

static_edges = {tuple(sorted(e)) for e in G.edges()}
dynamic_edges = set()
current_edges = []
final_path_edges = set()

pos = nx.spring_layout(G, seed=42)


def draw_graph(step=None, title_extra=''):
    plt.clf()
    title = f'Step {step}' if step is not None else 'Complete Path'
    plt.title(title + title_extra)
    colors = []
    for n in G.nodes():
        if n == source:
            colors.append('green')
        elif n == destination:
            colors.append('pink')
        else:
            colors.append('lightblue')
    nx.draw_networkx_nodes(G, pos, node_color=colors, node_size=500)

    static_draw = [e for e in static_edges if
                   e not in dynamic_edges and e not in final_path_edges and e not in current_edges]
    if static_draw:
        nx.draw_networkx_edges(G, pos, edgelist=static_draw, edge_color='gray', width=2)

    dynamic_draw = [e for e in dynamic_edges if e not in final_path_edges and e not in current_edges]
    if dynamic_draw:
        nx.draw_networkx_edges(G, pos, edgelist=dynamic_draw, edge_color='yellow', width=2)

    if current_edges:
        nx.draw_networkx_edges(G, pos, edgelist=current_edges, edge_color='red', width=4)

    final_draw = [e for e in final_path_edges if e not in current_edges]
    if final_draw:
        nx.draw_networkx_edges(G, pos, edgelist=final_draw, edge_color='green', width=4)

    nx.draw_networkx_labels(G, pos)
    plt.axis('off')
    plt.pause(0.001)


step = 1
messages = []
for nbr in nodes[source].neighbors:
    messages.append(('RREQ', nbr, source, source))
destination_found = False
dest_prev = None

# RREQ Process
while messages and not destination_found:
    print(f"Step {step}, Processing {len(messages)} messages")
    draw_graph(step)
    plt.pause(1)
    new_messages = []
    current_edges.clear()
    for msg in messages:
        msg_type, current, origin, prev = msg
        if msg_type == 'RREQ':
            if origin in nodes[current].received_rreq:
                continue
            nodes[current].received_rreq.add(origin)
            nodes[current].routing_table[origin] = prev
            edge = tuple(sorted((current, prev)))
            current_edges.append(edge)
            dynamic_edges.add(edge)
            if current == destination:
                destination_found = True
                dest_prev = prev
                break
            for nbr in nodes[current].neighbors:
                if nbr == prev:
                    continue
                new_messages.append(('RREQ', nbr, origin, current))
    step += 1
    messages = new_messages

if destination_found:
    print(f"Step {step}, Processing 0 messages (Destination {destination} reached)")
    draw_graph(step, ' (Destination reached)')
    plt.pause(3)
else:
    print("目的节点不可达.")
    draw_graph(step)
    plt.pause(3)
    plt.show()
    exit()

# RREP Process
if destination_found:
    messages = [('RREP', dest_prev, source, destination)]
    while messages:
        print(f"Step {step}, Processing {len(messages)} messages")
        draw_graph(step)
        plt.pause(1)
        new_messages = []
        current_edges.clear()
        for msg in messages:
            msg_type, current, target, prev = msg
            if msg_type == 'RREP':
                edge = tuple(sorted((current, prev)))
                current_edges.append(edge)
                dynamic_edges.add(edge)
                if current == target:
                    continue
                next_hop = nodes[current].routing_table.get(target)
                if next_hop is not None:
                    new_messages.append(('RREP', next_hop, target, current))
        step += 1
        messages = new_messages

    # Build and print path
    path = []
    cur = destination
    while cur != source:
        path.append(cur)
        prev = nodes[cur].routing_table[source]
        edge = tuple(sorted((cur, prev)))
        final_path_edges.add(edge)
        dynamic_edges.add(edge)
        cur = prev
    path.append(source)
    path.reverse()
    print("\n路由路径:", " -> ".join(map(str, path)))

    draw_graph(None, ' (Final Path)')
    plt.pause(3)
    plt.show()
